#include <iostream>
#include <winsock.h>
#include <fstream>
#include <time.h>
#include <string>
#include<queue>
#include<vector>

#include "reliableUdp.h"
#pragma comment(lib, "Ws2_32.lib")
using namespace std;

char buf[200000000];

SOCKET client;
SOCKADDR_IN serverAddr, clientAddr;

int base = 0;
int nextSeqNum = 0; //nextSeqNum = base;
vector<int> sentPktLen;
int timerBegin;

int sentNum = 0;//已经发送的数量
int sentRcvdNum = 0;//已经成功发送的数量
bool completed = false;

int cwnd = 1;
int ssthresh = 10;
int dupAckNum = 0;
int lastAck = -1;
int flag = 0; //当前状态——慢启动0，拥塞避免1

bool shakeHand() {
	while (true) {
		char shake_package[2];
		shake_package[1] = SHAKE_1;
		shake_package[0] = checksum(shake_package + 1, 1);
		sendto(client, shake_package, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
		cout << "Sent Shake 1.\n";
		//第一次握手
		int begin = clock();
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		bool lost = false;
		while (recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen) == SOCKET_ERROR) {
			if (clock() - begin > MAX_WAIT_TIME) {
				lost = true;
				cout << "Time limit for receiving Shake 2 exceeded.\n";
				break;
			}
		}
		if (!lost && checksum(recv, 2) == 0 && recv[1] == SHAKE_2) {
			cout << "Received Shake 2.\n";
			//第三次握手
			shake_package[1] = SHAKE_3;
			shake_package[0] = checksum(shake_package + 1, 1);
			sendto(client, shake_package, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
			cout << "Sent Shake 3.\n";
			break;
		}
	}
	while (true) {
		bool hasResentShake2 = true;
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		int begin = clock();
		while (recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen) == SOCKET_ERROR) {
			if (clock() - begin > MAX_WAVE_TIME) {
				hasResentShake2 = false;
				break;
			}
		}
		if (hasResentShake2 && checksum(recv, 2) == 0 && recv[1] == SHAKE_2) {
			cout << "Received Shake 2.\n";
			char send[2];
			send[1] = SHAKE_3;
			send[0] = checksum(send + 1, 1);
			sendto(client, send, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
			cout << "Sent Shake 3.\n";
			continue;
		}
		break;
	}
	cout << "Shaking hand completed.\n";
	return true;
}
bool waveHand() {
	while (true) {
		char send[2];
		send[1] = WAVE_1;
		send[0] = checksum(send + 1, 1);
		sendto(client, send, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
		cout << "Sent Wave 1.\n";
		int begin = clock();
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		bool lost = false;
		while (true) {
			int recvRtn = recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen);
			if (recvRtn > 0 && checksum(recv, 2) == 0) {
				if (recv[1] == WAVE_2) {
					cout << "Received Wave 2.\n";
					lost = false;
					break;
				}
				if (recv[1] == WAVE_3) {
					cout << "Received Wave 3.\n";
					lost = false;
					break;
				}
			}
			else {
				if (clock() - begin > MAX_WAIT_TIME) {
					lost = true;
					cout << "Time limit for receiving Wave 2 exceeded.\n";
					break;
				}
			}
		}
		if (!lost && checksum(recv, 2) == 0 && recv[1] == WAVE_2) {
			break;
		}
	}
	while (true) {
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		int begin = clock();
		while (recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen) == SOCKET_ERROR);
		if (checksum(recv, 2) != 0 || recv[1] != (char)WAVE_3) {
			continue;
		}
		cout << "Received Wave 3.\n";
		char send[2];
		send[1] = WAVE_4;
		send[0] = checksum(send + 1, 1);
		sendto(client, send, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
		cout << "Sent Wave 4.\n";
		break;
	}
	while (true) {
		bool hasResentWave3 = true;
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		int begin = clock();
		while (recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen) == SOCKET_ERROR) {
			if (clock() - begin > MAX_WAVE_TIME) {
				hasResentWave3 = false;
				break;
			}
		}
		if (hasResentWave3 && checksum(recv, 2) == 0 && recv[1] == WAVE_3) {
			cout << "Received Wave 3.\n";
			char send[2];
			send[1] = WAVE_4;
			send[0] = checksum(send + 1, 1);
			sendto(client, send, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
			cout << "Sent Wave 4.\n";
			continue;
		}
		break;
	}
	cout << "Waving hand completed.\n";
	return true;
}

int gbnSendPkt(char* pkt, int len, bool last, char*& formedPkt) {
	if (len > MAX_LEN || (!last && len != MAX_LEN)) {
		cout << "[Error] Wrong packet splitted size.\n";
		return -2;
	}
	if (nextSeqNum >= base + cwnd) {
		return -1;
	}
	int formedPktLen;
	if (!last) {
		formedPkt = new char[len + 3];
		formedPkt[1] = MIDDLE;
		formedPkt[2] = nextSeqNum;
		for (int i = 3; i < len + 3; i++) {
			formedPkt[i] = pkt[i - 3];
		}
		formedPkt[0] = checksum(formedPkt + 1, len + 2);
		formedPktLen = len + 3;
	}
	else {
		cout << "\n\nTHIS IS THE LAST!\n\n\n";
		formedPkt = new char[len + 5];
		formedPkt[1] = LAST;
		formedPkt[2] = nextSeqNum % 256;
		formedPkt[3] = len >> 8;
		formedPkt[4] = len & 0xFF;
		for (int i = 5; i < len + 5; i++) {
			formedPkt[i] = pkt[i - 5];
		}
		formedPkt[0] = checksum(formedPkt + 1, len + 4);
		formedPktLen = len + 5;
	}
	cout << "checksum = " << (int)((unsigned char)formedPkt[0]) << ": \n";
	//每发送一个分组nextseqnum就会增加，但不可以超过窗口长度
	sentPktLen.push_back(formedPktLen);
	sendto(client, formedPkt, sentPktLen[nextSeqNum], 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
	nextSeqNum++; //发送后，增加nextseqnum
	sentNum++;
	return base + cwnd - nextSeqNum;
}

DWORD WINAPI receiver(LPVOID lpParameter) {
	char recv[2];
	int serverAddrLen = sizeof(serverAddr);
	while (!completed) {
		while (true) {
			int recvRtn = recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen);
			if (recvRtn > 0) {
				if (checksum(recv, 2) == 0 && recv[1] != NAK) {
					int rcvNum = base;
					if ((int)((unsigned char)recv[1]) >= base) {
						base = (int)((unsigned char)recv[1]) + 1; //update base
					}
					if ((int)((unsigned char)recv[1]) == lastAck) {
						dupAckNum++;
					}
					else {
						//update cwnd
						switch (flag) {
						case 0: //slow start
							cwnd *= 2;
							break;
						case 1: //congestion avoidance
							cwnd++;
							break;
						default:
							cout << "[BUG] Unexpected flag??\n";
							break;
						}
					}
					if (cwnd > ssthresh && flag == 0) {
						flag = 1;
						cout << "[STATE]Transfer to congestion avoidance, for cwnd(" << cwnd << ") > ssthresh(" << ssthresh << ")\n";
					}
					cout << "Received ACK, ackNum = " << (int)((unsigned char)recv[1]) << ", base = " << base << endl;
					rcvNum = base - rcvNum;
					sentRcvdNum += rcvNum;
					lastAck = (int)((unsigned char)recv[1]);
					timerBegin = clock(); //reset timer
					break;
				}
				else {
					if ((unsigned char)(recv[1]) == NAK) {
						cout << "[Warning] Received NAK.\n";
						break;
					}
					else {
						cout << "[Error] Received something not ACK or NAK.\n";
						return false;
					}
				}
			}
		}
	}
}

bool gbnSender(char* file, int len) {
	completed = false;
	int pktNum = len / MAX_LEN + (len % MAX_LEN != 0);
	cout << "pktNum = " << pktNum << endl;
	char** formedPkts = new char* [pktNum];
	while (sentRcvdNum < pktNum) {
		if (sentNum < pktNum) {
			int sendRtn = gbnSendPkt(file + sentNum * MAX_LEN, sentNum == pktNum - 1 ? len - (pktNum - 1) * MAX_LEN : MAX_LEN, sentNum == pktNum - 1, formedPkts[nextSeqNum]);
			if (sendRtn == -2) {
				cout << "[Error] gbnSendPkt failure.\n";
				return false;
			}
			if (sendRtn >= 0) {
				cout << "Sent packet " << sentNum << ", seqNum = " << nextSeqNum - 1 << ", Len = " << (sentRcvdNum == pktNum - 1 ? len - (pktNum - 1) * MAX_LEN + 5 : MAX_LEN + 3) << ", WIN = " << sendRtn << ", cwnd = " << cwnd << endl;
			}
			if (base == nextSeqNum - 1) { //如果发送的是窗口里的第一个分组，那么需要启动timer
				//cout << "Start timer.\n";
				timerBegin = clock();
			}
		}
		//快速恢复
		if (clock() - timerBegin > MAX_WAIT_TIME || dupAckNum >= 3) {
			cout << "[Warning] Time limit or dupAckNum exceeded. Resend.\n";
			for (int i = base; i < nextSeqNum; i++) {
				sendto(client, formedPkts[i], sentPktLen[i], 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
				cout << "Resent packet " << i << ", seqNum = " << (int)((unsigned char)(formedPkts[i][2])) << ", Len = " << sentPktLen[i] << ", WIN = " << base + cwnd - i - 1 << ", cwnd = " << cwnd << endl;
			}
			if (dupAckNum >= 3) {
				ssthresh = cwnd / 2;
				cwnd = ssthresh + 3;
				flag = 1; //进入拥塞避免
				cout << "[STATE]Transfer to congestion avoidance, for dupAckNum >= 3...\n";
			}
			else {
				ssthresh = cwnd / 2;
				cwnd = 1;
				flag = 0; //进入慢启动
				cout << "[STATE]Transfer to slow start, for time limit exceeded...\n";
			}
			timerBegin = clock(); //重置计时器
			dupAckNum = 0;
		}
	}
	cout << "Sent successfully." << endl;
	completed = true;
	sentNum = 0;
	sentRcvdNum = 0;
	base = 0;
	nextSeqNum = 0;
	cwnd = 1;
	ssthresh = 10;
	dupAckNum = 0;
	lastAck = -1;
	flag = 0;
	sentPktLen.clear();
	for (int i = 0; i < pktNum; i++) {
		delete[] formedPkts[i];
	}
	delete[] formedPkts;
	formedPkts = nullptr;
	return true;
}

int main() {
	WSADATA wsadata;
	if (WSAStartup(MAKEWORD(2, 2), &wsadata)) {
		printf("error!");
		return 0;
	}

	serverAddr.sin_family = AF_INET;
	serverAddr.sin_port = htons(serverPort);
	serverAddr.sin_addr.s_addr = inet_addr("127.0.0.1");

	clientAddr.sin_family = AF_INET;
	clientAddr.sin_port = htons(clientPort);
	clientAddr.sin_addr.s_addr = inet_addr("127.0.0.1");

	client = socket(AF_INET, SOCK_DGRAM, 0);

	/*
	unsigned long ul = 1;
	if (ioctlsocket(client, FIONBIO, (unsigned long*)&ul) == SOCKET_ERROR) {
		cout << "Set failed.\n";
	}
	*/
	int time_out = 2000;
	setsockopt(client, SOL_SOCKET, SO_RCVTIMEO, (char*)&time_out, sizeof(time_out));

	if (client == INVALID_SOCKET) {
		printf("socket of client invalid!");
		closesocket(client);
		return 0;
	}

	if (bind(client, (sockaddr*)(&clientAddr), sizeof(clientAddr)) == SOCKET_ERROR) {
		printf("bind fail");
		closesocket(client);
		WSACleanup();
		return 0;
	}
	//握手
	cout << "Trying to establish connection...\n";
	shakeHand();
	cout << "Connection with server " << inet_ntoa(serverAddr.sin_addr) << ": " << ntohs(serverAddr.sin_port) << " established.\n===================================\n";
	while (true) {
		HANDLE hThread2 = CreateThread(NULL, NULL, &receiver, NULL, 0, NULL);
		SetThreadPriority(hThread2, THREAD_PRIORITY_HIGHEST);
		CloseHandle(hThread2);
		//传文件
		string fileName;
		cout << "Please input your filename: \n";
		cin >> fileName;
		gbnSender((char*)fileName.c_str(), fileName.length());
		//exit
		if (!strcmp("/exit", fileName.c_str())) {
			completed = true;
			break;
		}
		// 使用二进制方式 打开当前目录下的文件
		string fileDir = "C:\\Users\\16116\\Desktop\\rt\\file\\";
		string filePath = fileDir + fileName;
		ifstream in(filePath.c_str(), ifstream::binary);
		int len = 0;
		if (!in) {
			cout << "can't open the file! Please retry\n";
			return 0;
		}
		// 文件读取到buffer
		BYTE t = in.get();
		while (in) {
			buf[len++] = t;
			t = in.get();
		}
		in.close();
		gbnSender(buf, len);
	}
	waveHand();
	closesocket(client);
	cout << "Connection Aborted.\n";
	WSACleanup();
	cout << "WSA cleaning up finished.\n";
	return 0;
}